#!/usr/bin/python

#  Copyright (C) 2012 Hai Bison
#
#  See the file LICENSE at the root directory of this project for copying
#  permission.

License = '''
###########
###
##      Lines
#
#   Version     :   1.8
#   Release     :   November 07th, 2012
#   Requires    :   Python >= 2.6.8
#   License     :   MIT License

Copyright (C) 2012 Hai Bison

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
'''

import argparse
import datetime
import os
import os.path
import re
import signal
import sys
import threading
import time
import traceback

class MaxFileSizeAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        if not isinstance(values, str) or not re.match(r'[0-9]+[KM]', values):
            raise Exception('Invalid option for "%s"' % option_string)
        factor = 1024 if re.match('.+K', values) else 1024 * 1024
        values = re.sub(r'[^0-9]', '', values)
        if not values: raise Exception('Invalid option for "%s"' % option_string)
        setattr(namespace, self.dest, int(values) * factor)
    #/__call__()
#/MaxFileSizeAction()

parser = argparse.ArgumentParser(
    prog='lines',
    description='Counts all lines of a project.')
parser.add_argument('root', metavar='root', nargs='?', default='', type=str,
    help='can be either a root directory of project, or a source file;')
parser.add_argument('-R', '--recursive', action='store_true',
    help='''if ROOT is a directory, this flag lets the app browse and handle
         all of its files recursively; default is false;''')
parser.add_argument('-r', '--regex', type=str, default=None,
    help='regular expression to filter files; default is None;')
parser.add_argument('-rxd', '--regex-exclude-dirs', type=str, default=None,
    dest='rxd',
    help='''regular expression to exclude directories; this flag is useful
        in case you want to exclude directories of version control system;
        e.g: ".hg", ".git"...''')
parser.add_argument('-max', default=1*1024*1024, # 1 MiB
    action=MaxFileSizeAction,
    help='''set maximum file size which is allowed to be handled; the form
        is "size" + "K|M"; e.g: 5K, 5M... default is "1M" (1 MiB);''')
parser.add_argument('-D', '--debug', action='store_true',
    help='print debug information;')
parser.add_argument('-v', '--version', action='store_true',
    help='print version + license information and exit;')

args = parser.parse_args()

class SpeedTracker():

    class CheckPoint():
        def __init__(self, speed, tick = 0):
            self.tick = tick if tick > 0 else time.time()
            self.speed = speed
        #/__init__()
    #/CheckPoint

    def __init__(self, max_checkpoints = 500, max_period = 10):
        self.max_checkpoints = max_checkpoints
        self.max_period = max_period
        self.checkpoints = []
    #/__init__()

    def add(self, speed):
        self.checkpoints.append(self.CheckPoint(speed = speed))
        if len(self.checkpoints) > self.max_checkpoints:
            del self.checkpoints[0]
    #/add()

    def clear(self):
        del self.checkpoints[:]
    #/clear()

    def instantaneous_speed(self):
        tick = time.time()
        total_speed = 0
        for i in range(len(self.checkpoints) - 1, -1, -1):
            cp = self.checkpoints[i]
            if tick - cp.tick <= self.max_period:
                total_speed += cp.speed
            else:
                del self.checkpoints[:i+1]
                break
        #/for

        return 0 if len(self.checkpoints) == 0 else total_speed / len(self.checkpoints)
    #/instantaneous_speed()
#/SpeedTracker

# UTILITIES

NAN = 1e-9
BLOCK_SIZE = 2 * 1024 * 1024 # 2 MiB
ctrl_c = False

def count_lines(f, speed_tracker):
    size = os.path.getsize(f)

    if args.debug:
        print(' ? "{:s}" ({:,} bytes)'.format(f, size))

    if size <= 0: return 0
    if size > args.max: return -1

    pattern = re.compile(b'\n')
    count = 1 # if file size > 0, at least there is always one line
    ctrl_c = False

    signal.alarm(5)
    f = open(f, mode='r+b', buffering=BLOCK_SIZE)
    signal.alarm(0)
    try:
        while not ctrl_c:
            tick = time.time()
            data = f.read(BLOCK_SIZE)
            if data:
                iterator = re.finditer(pattern, data)
                for match in iterator: count += 1

                tick = max(time.time() - tick, NAN)
                speed_tracker.add(len(data) / tick)
            else: break
    finally: f.close()

    return count
#/count_lines()

# THREADS

class BaseThread(threading.Thread):

    terminated = False
    final_statement = None

    def terminate(self, final_statement = None):
        self.terminated = True
        if final_statement != None: self.final_statement = final_statement
#/BaseThread

class Reporter(BaseThread):

    def __init__(self, handler):
        super(Reporter, self).__init__()
        self.handler = handler

    def run(self):
        self.start_time = time.time()
        while not self.terminated:
            self.handler()
            time.sleep(0.5)
    # run()
#/Reporter

def size_to_str(size):
    import math
    if math.isnan(size): return 'NaN'
    if size <= 0: return '0 B'

    units = ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi', 'Yi']
    b_size = 1024

    import math
    digit_groups = int(min(math.log10(size) / math.log10(b_size), len(units) - 1))

    size /= b_size**digit_groups
    size_s = '{:,.{p}f}'.format(size, p=(0 if digit_groups == 0 else 2))
    return '%s %sB' % (size_s, units[digit_groups])
#/size_to_str()

class Counter(BaseThread):

    start_time = time.time()
    line_count = 0
    file_count = 0
    total_file_count = 0
    total_size = 0

    reporter = None
    speed_tracker = SpeedTracker()

    _header_pattern = ' {:19s} {:11s} {:11s} {:16s} {:s}'

    def print_headers(self):
        print(self._header_pattern.format('[Lines]', '[Files]', '[All files]', '[Speed]', '[Time]'))
    #/print_header()

    def print_stats(self, summary=False):
        if summary:
            speed = size_to_str(self.total_size / max(time.time() - self.start_time, NAN))
        else:
            speed = size_to_str(self.speed_tracker.instantaneous_speed())
        speed += '/s'

        s = '\b' * 80 + self._header_pattern.format(
            '{:,}'.format(self.line_count), '{:,}'.format(self.file_count),
            '{:,}'.format(self.total_file_count), speed,
            str(datetime.timedelta(seconds = time.time() - self.start_time)))

        if summary: print(s + '\b' * (80 - len(s)))
        else:
            sys.stdout.write(s)
            sys.stdout.flush()

        return len(s)
    #/print_stats()

    def run(self):
        self.print_headers()

        self.reporter = Reporter(self.print_stats)
        self.reporter.daemon = True
        self.reporter.start()

        if os.path.isfile(args.root):
            self.total_file_count += 1
            try:
                c = count_lines(args.root, self.speed_tracker)
                if c >= 0:
                    self.total_size += os.path.getsize(args.root)
                    self.line_count += c
                    self.file_count = 1
            except: pass
        elif os.path.isdir(args.root):
            for path, dirs, files in os.walk(args.root, True):
                if self.terminated: break

                for f in files:
                    if self.terminated: break

                    if args.regex == None or re.match(args.regex, f):
                        # ignore excluded dirs
                        if args.rxd != None:
                            found = False
                            p = path
                            while p and not os.path.samefile(args.root, p):
                                if re.match(args.rxd, os.path.basename(p)):
                                    found = True; break
                                p = os.path.dirname(p)
                            if found: continue

                        self.total_file_count += 1

                        try:
                            c = count_lines(os.path.join(path, f), self.speed_tracker)
                        except KeyboardInterrupt:
                            self.terminate('! User interrupted...')
                        except Exception as e:
                            if args.debug: print(' ! Error: "%s"' % e)
                            continue

                        if c >= 0:
                            self.total_size += os.path.getsize(os.path.join(path, f))
                            self.line_count += c
                            self.file_count += 1
                if self.terminated or not args.recursive: break
        else:
            self.reporter.terminate('! ROOT is invalid.')

        if not self.reporter.terminated: self.reporter.terminate()

        if self.final_statement != None:
            length = self.print_time()
            print('\n' + self.final_statement)
        else:
            self.print_stats(True)
    #/run()
#/Counter()

# MAIN

def main():
    if args.version:
        import sys
        print(License)
        sys.exit()

    if args.debug:
        print(' ? Max filesize = {:,} bytes'.format(args.max))

    # handle IO blocking
    def err(signum, frame):
        signal.alarm(0)
        raise IOError('Error accessing some file...')
    signal.signal(signal.SIGALRM, err)

    counter = Counter()
    counter.daemon = True
    counter.start()
    while counter.is_alive():
        try: counter.join(10)
        except (KeyboardInterrupt, SystemExit):
            ctrl_c = True
            counter.terminate('! User interrupted...')
            break
#/main()

if __name__ == '__main__':
    main()
